package com.sso.oauth.utils;

import com.alibaba.fastjson.JSON;
import com.microsoft.aad.adal4j.AuthenticationContext;
import com.microsoft.aad.adal4j.AuthenticationResult;
import com.microsoft.aad.adal4j.ClientCredential;
import com.nimbusds.jwt.JWTParser;
import com.nimbusds.oauth2.sdk.AuthorizationCode;
import com.nimbusds.openid.connect.sdk.AuthenticationErrorResponse;
import com.nimbusds.openid.connect.sdk.AuthenticationResponse;
import com.nimbusds.openid.connect.sdk.AuthenticationResponseParser;
import com.nimbusds.openid.connect.sdk.AuthenticationSuccessResponse;
import com.sso.oauth.config.SSOConfigProperties;
import com.sso.oauth.constants.AuthParameterNames;
import com.sso.oauth.vo.StateData;
import lombok.extern.log4j.Log4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.naming.ServiceUnavailableException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.HttpURLConnection;
import java.net.URI;
import java.net.URL;
import java.net.URLEncoder;
import java.text.ParseException;
import java.util.*;
import java.util.concurrent.*;

/**
 * Created by Derry Luo on 2019/5/10.
 */
@Log4j
@Component
public class SessionHelper {

    public static final Integer STATE_TTL = 3600;
    public static final String FAILED_TO_VALIDATE_MESSAGE = "Failed to validate data received from Authorization service - ";

    @Autowired
    private SSOConfigProperties ssoConfigProperties;


    // 校验用户是否已登录
    public boolean isAuthenticated(HttpServletRequest httpRequest) {
        return httpRequest.getSession().getAttribute(AuthParameterNames.USER_INFO_SESSION_KEY) != null;
    }

    /**
     * 判断accessToken是否已过期
     *
     * @param httpRequest
     * @return
     */
    public boolean isAuthDataExpired(HttpServletRequest httpRequest) {
        log.info("【Validate AccessToken expireDate】");
        AuthenticationResult authData = getAuthSessionObject(httpRequest);
        if (authData == null) {
            return false;
        }
        return authData.getExpiresOnDate().before(new Date()) ? true : false;
    }

    private AuthenticationResult getAuthSessionObject(
            HttpServletRequest request) {
        return (AuthenticationResult) request.getSession().getAttribute(
                AuthParameterNames.PRINCIPAL_SESSION_NAME);
    }

    /**
     * 刷新accessToken
     *
     * @param httpRequest
     * @throws Throwable
     */
    public void updateAuthDataUsingRefreshToken(HttpServletRequest httpRequest) throws Throwable {
        log.info("【Refresh AccessToken by RefreshToken】");
        AuthenticationResult authData = getAccessTokenFromRefreshToken(getAuthSessionObject(httpRequest).getRefreshToken());
        // 把获取AccessToken时获取的完整信息存放起来
        setSessionPrincipal(httpRequest, authData);
    }

    /**
     * 通过refreshToken刷新accessToken
     *
     * @param refreshToken
     * @return
     * @throws Throwable
     */
    private AuthenticationResult getAccessTokenFromRefreshToken(
            String refreshToken) throws Throwable {
        AuthenticationContext context;
        AuthenticationResult result = null;
        ExecutorService service = null;
        try {
            service = Executors.newFixedThreadPool(1);
            context = new AuthenticationContext(ssoConfigProperties.getAuthority() + ssoConfigProperties.getTenant() + "/", true,
                    service);
            Future<AuthenticationResult> future = context
                    .acquireTokenByRefreshToken(refreshToken, new ClientCredential(ssoConfigProperties.getClientId(), ssoConfigProperties.getSecretKey()), null, null);
            result = future.get();
        } catch (ExecutionException e) {
            throw e.getCause();
        } finally {
            service.shutdown();
        }

        if (result == null) {
            throw new ServiceUnavailableException("authentication result was null");
        }
        return result;
    }

    private void setSessionPrincipal(HttpServletRequest httpRequest, AuthenticationResult result) {
        log.info("【Save Original Auth Data into Session】");
        httpRequest.getSession().setAttribute(AuthParameterNames.PRINCIPAL_SESSION_NAME, result);
    }

    /**
     * 判断Request请求信息中是否包含了身份信息或错误信息
     *
     * @param httpRequest
     * @return
     */
    public boolean containsAuthenticationData(HttpServletRequest httpRequest) {

        Map<String, String[]> map = httpRequest.getParameterMap();
        log.info("【Validate if Request Include AuthData】param = " + JSON.toJSONString(map));
        return httpRequest.getMethod().equalsIgnoreCase("POST") && (map.containsKey(
                AuthParameterNames.ERROR)
                || map.containsKey(
                AuthParameterNames.ID_TOKEN) || map.containsKey(AuthParameterNames.CODE));
    }

    /**
     * 从Request中拿出授权信息
     * <p>获取authCode</p>
     * <p>根据authCode获取accessToken、refreshToken等</p>
     * <p>将accessToken、refreshToken等存入session</p>
     *
     * @param httpRequest
     * @param currentUri
     * @param fullUrl
     * @throws Throwable
     */
    public void processAuthenticationData(HttpServletRequest httpRequest, String currentUri, String fullUrl)
            throws Throwable {

        HashMap<String, String> params = new HashMap<>();
        for (String key : httpRequest.getParameterMap().keySet()) {
            params.put(key, httpRequest.getParameterMap().get(key)[0]);
        }
        log.info("【Get Auth Data from Request】params = " + JSON.toJSONString(params));
        // validate that state in response equals to state in request
        StateData stateData = validateState(httpRequest.getSession(), params.get(AuthParameterNames.STATE));

        AuthenticationResponse authResponse = AuthenticationResponseParser.parse(new URI(fullUrl), params);
        if (isAuthenticationSuccessful(authResponse)) {
            log.info("【Authorized success, begin to analysis data】");
            AuthenticationSuccessResponse oidcResponse = (AuthenticationSuccessResponse) authResponse;
            // validate that OIDC Auth Response matches Code Flow (contains only requested artifacts)
            validateAuthRespMatchesCodeFlow(oidcResponse);

            AuthenticationResult authData = getAccessToken(oidcResponse.getAuthorizationCode(), currentUri);
            // validate nonce to prevent reply attacks (code maybe substituted to one with broader access)
            validateNonce(stateData, getClaimValueFromIdToken(authData.getIdToken(), "nonce"));

            // 将accessToken等信息保存到session中
            setSessionPrincipal(httpRequest, authData);
        } else {
            AuthenticationErrorResponse oidcResponse = (AuthenticationErrorResponse) authResponse;
            throw new Exception(String.format("Request for auth code failed: %s - %s",
                    oidcResponse.getErrorObject().getCode(),
                    oidcResponse.getErrorObject().getDescription()));
        }
    }

    private boolean isAuthenticationSuccessful(
            AuthenticationResponse authResponse) {
        return authResponse instanceof AuthenticationSuccessResponse;
    }

    /**
     * 系统内的隐式登录
     * <p>根据accessToken获取用户信息</p>
     * <p>根据用户信息做一次系统内部的登录</p>
     * <p>将系统内部用户登录信息存入session</p>
     *
     * @param httpServletRequest
     */
    public void processUserLogin(HttpServletRequest httpServletRequest, AuthenticationResult authenticationResult) {
        this.processUserLogin(httpServletRequest, authenticationResult.getAccessToken());
    }

    /**
     * 系统内的隐式登录
     * <p>根据accessToken获取用户信息</p>
     * <p>根据用户信息做一次系统内部的登录</p>
     * <p>将系统内部用户登录信息存入session</p>
     *
     * @param httpServletRequest
     */
    public void processUserLogin(HttpServletRequest httpServletRequest, String accessToken) {
        log.info("【Process User Login】accessToken = " + accessToken);
        String userJsonInfo = null;
        try {
            userJsonInfo = getUsernameFromGraph(accessToken);
        } catch (Exception e) {
            throw new RuntimeException("【Fail to Get User Info from Graph】", e);
        }
        log.info("【User info from Graph】userJsonInfo = " + userJsonInfo);
        // TODO 系统内部的登录
        httpServletRequest.getSession().setAttribute(AuthParameterNames.USER_INFO_SESSION_KEY, userJsonInfo);
    }

    /**
     * 使用accessToken从Azure获取用户的信息
     *
     * @param accessToken
     * @return
     * @throws Exception
     */
    private String getUsernameFromGraph(String accessToken) throws Exception {

        log.info("【Request username From Graph】accessToken = " + accessToken);
        URL url = new URL(ssoConfigProperties.getGraphUrl());
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();

        conn.setRequestMethod("GET");
        conn.setRequestProperty("Authorization", "Bearer " + accessToken);
        conn.setRequestProperty("Accept", "application/json");
        int httpResponseCode = conn.getResponseCode();

        String goodRespStr = HttpClientHelper.getResponseStringFromConn(conn, true);

        // log.info("goodRespStr ->" + goodRespStr);
        return goodRespStr;
    }

    private void validateNonce(StateData stateData, String nonce) throws Exception {
        log.info("【validate Nonce】");
        if (StringUtils.isEmpty(nonce) || !nonce.equals(stateData.getNonce())) {
            throw new Exception(FAILED_TO_VALIDATE_MESSAGE + "could not validate nonce");
        }
    }

    private String getClaimValueFromIdToken(String idToken, String claimKey) throws ParseException {
        return (String) JWTParser.parse(idToken).getJWTClaimsSet().getClaim(claimKey);
    }

    private AuthenticationResult getAccessToken(
            AuthorizationCode authorizationCode, String currentUri)
            throws Throwable {
        String authCode = authorizationCode.getValue();
        ClientCredential credential = new ClientCredential(ssoConfigProperties.getClientId(), ssoConfigProperties.getSecretKey());
        AuthenticationContext context;
        AuthenticationResult result = null;
        ExecutorService service = null;
        try {
            service = Executors.newFixedThreadPool(1);
            context = new AuthenticationContext(ssoConfigProperties.getAuthority() + ssoConfigProperties.getTenant() + "/", true,
                    service);
            log.info("【Get AccessToken by authCode】");
            Future<AuthenticationResult> future = context
                    .acquireTokenByAuthorizationCode(authCode, new URI(
                            currentUri), credential, null);
            result = future.get();
        } catch (ExecutionException e) {
            throw e.getCause();
        } finally {
            service.shutdown();
        }

        if (result == null) {
            throw new ServiceUnavailableException("authentication result was null");
        }
        return result;
    }

    private void validateAuthRespMatchesCodeFlow(AuthenticationSuccessResponse oidcResponse) throws Exception {
        if (oidcResponse.getIDToken() != null || oidcResponse.getAccessToken() != null ||
                oidcResponse.getAuthorizationCode() == null) {
            throw new Exception(FAILED_TO_VALIDATE_MESSAGE + "unexpected set of artifacts received");
        }
    }

    /**
     * make sure that state is stored in the session,
     * delete it from session - should be used only once
     *
     * @param session
     * @param state
     * @throws Exception
     */
    private StateData validateState(HttpSession session, String state) throws Exception {
        if (StringUtils.isNotEmpty(state)) {
            StateData stateDataInSession = removeStateFromSession(session, state);
            if (stateDataInSession != null) {
                return stateDataInSession;
            }
        }
        throw new Exception(FAILED_TO_VALIDATE_MESSAGE + "could not validate state");
    }

    @SuppressWarnings("unchecked")
    private StateData removeStateFromSession(HttpSession session, String state) {
        Map<String, StateData> states = (Map<String, StateData>) session.getAttribute(AuthParameterNames.STATES);
        if (states != null) {
            eliminateExpiredStates(states);
            StateData stateData = states.get(state);
            if (stateData != null) {
                states.remove(state);
                return stateData;
            }
        }
        return null;
    }

    private void eliminateExpiredStates(Map<String, StateData> map) {
        Iterator<Map.Entry<String, StateData>> it = map.entrySet().iterator();

        Date currTime = new Date();
        while (it.hasNext()) {
            Map.Entry<String, StateData> entry = it.next();
            long diffInSeconds = TimeUnit.MILLISECONDS.
                    toSeconds(currTime.getTime() - entry.getValue().getExpirationDate().getTime());

            if (diffInSeconds > STATE_TTL) {
                it.remove();
            }
        }
    }

    public void sendAuthRedirect(HttpServletRequest httpRequest, HttpServletResponse httpResponse) throws IOException {
        httpResponse.setStatus(302);

        // use state parameter to validate response from Authorization server
        String state = UUID.randomUUID().toString();

        // use nonce parameter to validate idToken
        String nonce = UUID.randomUUID().toString();

        storeStateInSession(httpRequest.getSession(), state, nonce);


//        String currentUri = httpRequest.getRequestURL().toString();
        StringBuffer url = httpRequest.getRequestURL();
        String currentUri = url.delete(url.length() - httpRequest.getRequestURI().length(), url.length()).append(httpRequest.getSession().getServletContext().getContextPath()).append("/").toString();
        currentUri = currentUri.replaceAll("http:", "https:") + httpRequest.getServletPath() + "?" + httpRequest.getQueryString();
        log.info("【User Request Uri = " + currentUri + "】");

        httpRequest.getSession().setAttribute(AuthParameterNames.GOTO_PATH, currentUri);

        httpResponse.sendRedirect(getRedirectUrl(ssoConfigProperties.getReplyPath(), state, nonce));
    }

    private String getRedirectUrl(String currentUri, String state, String nonce)
            throws UnsupportedEncodingException {
        log.info("【Redirect to Login Page】authority = " + ssoConfigProperties.getAuthority() + " , tenant = " + ssoConfigProperties.getTenant() + " , clientId = " + ssoConfigProperties.getClientId());
        String redirectUrl = ssoConfigProperties.getAuthority()
                + ssoConfigProperties.getTenant()
                + "/oauth2/authorize?response_type=code&scope=directory.read.all&response_mode=form_post&redirect_uri="
                + URLEncoder.encode(currentUri, "UTF-8") + "&client_id="
                + ssoConfigProperties.getClientId() + "&resource=https%3a%2f%2fgraph.microsoft.com"
                + "&state=" + state
                + "&nonce=" + nonce;
        log.info("【Redirect URL】URL = " + redirectUrl);
        return redirectUrl;
    }

    @SuppressWarnings("unchecked")
    private void storeStateInSession(HttpSession session, String state, String nonce) {
        if (session.getAttribute(AuthParameterNames.STATES) == null) {
            session.setAttribute(AuthParameterNames.STATES, new HashMap<String, StateData>());
        }
        ((Map<String, StateData>) session.getAttribute(AuthParameterNames.STATES)).put(state, new StateData(nonce, new Date()));
    }

    public void removePrincipalFromSession(HttpServletRequest httpRequest) {
        log.info("【Remove Principal From Session】");
        httpRequest.getSession().removeAttribute(AuthParameterNames.PRINCIPAL_SESSION_NAME);
    }

    /**
     * 判断请求信息中是否包含令牌：accessToken
     *
     * @param httpRequest
     * @return
     */
    public boolean containsAccessToken(HttpServletRequest httpRequest) {
        String accessToken = httpRequest.getParameter(AuthParameterNames.ACCESS_TOKEN);
        return StringUtils.isNotBlank(accessToken);
    }

    /**
     * 注销登录
     *
     * @param servletRequest
     */
    public void logout(HttpServletRequest servletRequest, HttpServletResponse servletResponse) throws IOException {
        log.info("【User Request Logout】");
        this.sendLogoutRedirect(servletRequest, servletResponse);
    }

    public void sendLogoutRedirect(HttpServletRequest httpRequest, HttpServletResponse httpResponse) throws IOException {
        httpResponse.setStatus(302);

        // https://login.microsoftonline.com/{tenant}/oauth2/v2.0/logout?post_logout_redirect_uri=https://localhost/myapp/
        StringBuilder builder = new StringBuilder(ssoConfigProperties.getAuthority());
        builder.append(ssoConfigProperties.getTenant()).append("/oauth2/v2.0/logout?post_logout_redirect_uri=")
                .append(URLEncoder.encode(ssoConfigProperties.getLogoutPath(), "UTF-8"));


        log.info("【Logout Redirect URL】URL = " + builder.toString());

        httpResponse.sendRedirect(builder.toString());
    }

    /**
     * 清除用户登录的session信息
     *
     * @param request
     */
    public void removeLoginSessionInfo(HttpServletRequest request) {
        HttpSession session = request.getSession();
        session.removeAttribute(AuthParameterNames.PRINCIPAL_SESSION_NAME);
        session.removeAttribute(AuthParameterNames.USER_INFO_SESSION_KEY);
    }
}
